﻿#pragma once

#include <llhttp.h>
#include <net.hpp>

#include <list>
#include <map>
#include <optional>
#include <regex>

namespace leak::net::http {

using Headers                         = std::multimap<std::string, std::string>;
constexpr char name_value_separator[] = { ':', ' ' };
constexpr char crlf[]                 = { '\r', '\n' };

struct Request {
    uint8_t     major_version;
    uint8_t     minor_version;
    std::string method;
    std::string path;
    Headers     headers;
    std::string body;
    bool        keepalive;
};

struct Response {
    uint8_t     major_version;
    uint8_t     minor_version;
    int         status;
    std::string status_string;
    Headers     headers;
    std::string body;
};

class Parser {
public:
    struct Header {
        std::string name;
        std::string value;
    };

    Parser() {
        settings_ = {
            [](llhttp_t* parser) {
                return ((Parser*)parser->data)->OnMessageBegin();
            },
            [](llhttp_t* parser, const char* data, size_t len) {
                return ((Parser*)parser->data)->OnUrl(data, len);
            },
            [](llhttp_t* parser, const char* data, size_t len) {
                return ((Parser*)parser->data)->OnStatus(data, len);
            },
            [](llhttp_t* parser, const char* data, size_t len) {
                return ((Parser*)parser->data)->OnHeaderField(data, len);
            },
            [](llhttp_t* parser, const char* data, size_t len) {
                return ((Parser*)parser->data)->OnHeaderValue(data, len);
            },
            [](llhttp_t* parser) {
                return ((Parser*)parser->data)->OnHeaderComplete();
            },
            [](llhttp_t* parser, const char* data, size_t len) {
                return ((Parser*)parser->data)->OnBody(data, len);
            },
            [](llhttp_t* parser) {
                return ((Parser*)parser->data)->OnMessageComplete();
            }
        };
        llhttp_init(&parser_, HTTP_BOTH, &settings_);
        parser_.data = this;
    }

    bool Parse(std::string_view data) { return llhttp_execute(&parser_, data.data(), (int)data.length()) == HPE_OK; }
    bool IsFinish() { return parser_.finish == 0 || parser_.finish == 1; }
    void Reset() { llhttp_reset(&parser_); }

    virtual int OnMessageBegin() { return 0; }
    virtual int OnUrl(const char* data, size_t len) { return 0; }
    virtual int OnStatus(const char* data, size_t len) { return 0; }
    virtual int OnHeaderField(const char* data, size_t len) { return 0; }
    virtual int OnHeaderValue(const char* data, size_t len) { return 0; }
    virtual int OnHeaderComplete() { return 0; }
    virtual int OnBody(const char* data, size_t len) { return 0; }
    virtual int OnMessageComplete() { return 0; }

protected:
    llhttp_t parser_;
    Header   header_;
    bool     write_v_;

private:
    llhttp_settings_t settings_;
};

class RequestParser : public Parser {
public:
    RequestParser() = default;

    int OnMessageBegin() override {
        write_v_ = false;
        header_.name.clear();
        header_.value.clear();
        request_.method.clear();
        request_.path.clear();
        request_.body.clear();
        request_.headers.clear();
        return 0;
    }

    int OnUrl(const char* data, size_t len) override {
        request_.path.append(data, len);
        return 0;
    }
    int OnHeaderField(const char* data, size_t len) override {
        if (write_v_) {
            request_.headers.emplace(std::make_pair(std::move(header_.name), std::move(header_.value)));
        }
        header_.name.append(data, len);
        write_v_ = false;
        return 0;
    }

    int OnHeaderValue(const char* data, size_t len) override {
        write_v_ = true;
        header_.value.append(data, len);
        return 0;
    }

    int OnHeaderComplete() override {
        request_.major_version = parser_.http_major;
        request_.minor_version = parser_.http_minor;
        request_.keepalive     = llhttp_should_keep_alive(&parser_);
        request_.method        = llhttp_method_name((llhttp_method)parser_.method);

        return 0;
    }

    int OnBody(const char* data, size_t len) override {
        request_.body.append(data, len);
        return 0;
    }

    Request& request() {
        return request_;
    }

private:
    Request request_;
};

class ResponseParser : public Parser {
public:
    ResponseParser() = default;

    int OnMessageBegin() override {
        write_v_ = false;
        response_.status_string.clear();
        response_.body.clear();
        response_.headers.clear();
        return 0;
    }

    int OnStatus(const char* data, size_t len) override {
        response_.status_string.append(data, len);
        return 0;
    }

    int OnHeaderField(const char* data, size_t len) override {
        if (write_v_) {
            response_.headers.emplace(std::make_pair(std::move(header_.name), std::move(header_.value)));
        }
        header_.name.append(data, len);
        write_v_ = false;
        return 0;
    }

    int OnHeaderValue(const char* data, size_t len) override {
        write_v_ = true;
        header_.value.append(data, len);
        return 0;
    }

    int OnHeaderComplete() override {
        if (write_v_) {
            response_.headers.emplace(std::make_pair(std::move(header_.name), std::move(header_.value)));
        }

        header_.name.clear();
        header_.value.clear();
        response_.major_version = parser_.http_major;
        response_.minor_version = parser_.http_minor;
        response_.status        = parser_.status_code;
        return 0;
    }

    int OnBody(const char* data, size_t len) override {
        response_.body.append(data, len);
        return 0;
    }

    Response& response() { return response_; }

private:
    Response response_;
};

// template <typename Socket>
class ResponseWriter : public Response, public std::enable_shared_from_this<ResponseWriter> {
public:
    ResponseWriter(std::shared_ptr<tcp::Socket> socket, bool close = false)
      : socket_(socket)
      , close_(close) {
        status        = 200;
        status_string = "OK";
    }

    void Write() {
        auto socket = socket_.lock();
        if (!socket) return;

        BuildHeader();
        socket->PrepareWrite("HTTP/1.1 " + std::to_string(status) + " " + status_string + http::crlf);

        for (auto& header : headers) {
            socket->PrepareWrite(header.first + ": " + header.second + http::crlf);
        }
        socket->PrepareWrite(http::crlf);
        socket->PrepareWrite(body);
        socket->Flush(close_);
    }

    void                         set_close(bool close) { close_ = close; }
    std::shared_ptr<tcp::Socket> socket() { return socket_.lock(); }

private:
    void BuildHeader() {
        if (body.size() > 0) {
            // Content-Type
            if (headers.count("Content-Type") == 0) {
                headers.emplace("Content-Type", "text/plain");
            }

            // Content-Length
            if (headers.count("Content-Length") == 0) {
                headers.emplace("Content-Length", std::to_string(body.size()));
            }
        }
    }

    std::weak_ptr<tcp::Socket> socket_;
    bool                       close_;
};

class Router {
    using Handler = std::function<void(const Request&, std::shared_ptr<ResponseWriter>)>;

public:
    virtual ~Router() = default;

    void Route(const Request& req, std::shared_ptr<ResponseWriter> rsp) {
        auto& headers = req.headers;

        auto it = handlers_.find(req.path);

        if (it == handlers_.end()) {
            rsp->status        = 404;
            rsp->status_string = "Not Found";
            rsp->body          = "Not Found";
            rsp->Write();
            return;
        }

        // keep-alive
        auto iter = headers.find("Connection");
        if (iter == headers.end()) {
            rsp->set_close(true);
        } else {
            rsp->set_close(iter->second == "close");
        }
        std::invoke(it->second, req, rsp);
    }

    void RegisterHandler(std::string path, Handler handler) {
        handlers_.emplace(std::make_pair(std::move(path), std::move(handler)));
    }

private:
    std::unordered_map<std::string, Handler> handlers_;
};

inline void NewHttpServer(utils::ContextPool pool, std::string_view host, uint16_t port, std::shared_ptr<Router> router) {
    assert(router);
    tcp::Server::New(pool->Next())
        ->set_context_pool(pool)
        ->Listen(host, port)
        ->Accept([router](std::shared_ptr<net::tcp::Socket> socket) {
            auto parser = std::make_shared<RequestParser>();
            socket
                ->set_onerror([socket](std::string_view error) {
                    std::cout << "Http socket err: " << error << std::endl;
                })
                ->StartRead([router, parser](std::shared_ptr<net::tcp::Socket> socket, std::string_view data) {
                    if (!parser->Parse(data)) {
                        auto response = std::make_shared<http::ResponseWriter>(socket, true);

                        response->status        = 400;
                        response->status_string = "Bad Request";
                        response->Write();
                        return -1;
                    } else if (parser->IsFinish()) {
                        router->Route(parser->request(), std::make_shared<http::ResponseWriter>(socket, true));
                        parser->Reset();
                        return -1;
                    }
                    return 0;
                });
        });
}

class Client : public std::enable_shared_from_this<Client> {
    using ResponseHandler = std::function<void(const Response&)>;
    using ErrorHandler    = std::function<void(std::string_view)>;
    struct Task {
        ResponseHandler          response_handler;
        ErrorHandler             error_handler;
        std::shared_ptr<Request> request;
        bool                     complete;
    };

public:
    static std::shared_ptr<Client> New(asio::io_context& ctx, std::string address) {
        return std::shared_ptr<Client>(new Client(ctx, std::move(address)));
    }

    void DoRequest(std::shared_ptr<Request> request, ResponseHandler response_handler, ErrorHandler error_handler = nullptr) {
        BuildHeaders(request);

        tasks_.emplace_back(std::make_shared<Task>(Task{
            std::move(response_handler),
            std::move(error_handler),
            request,
            false }));

        DoRequest();
    }

    void Close() { socket_.reset(); }

    Client* set_keep_alive(bool on) {
        keep_alive_ = on;
        return this;
    }

private:
    Client(asio::io_context& ctx, std::string address)
      : context_(ctx)
      , is_ssl_(false)
      , keep_alive_(false)
      , intask_(false) {
        if (std::regex_search(address, std::regex("https://"))) {
            is_ssl_  = true;
            address  = address.substr(8);
            auto pos = address.find(':');
            host_    = address.substr(0, pos);
            pos == std::string::npos ?
                   port_ = 443 :
                   port_ = std::stoi(address.substr(address.find(':') + 1));
        } else if (std::regex_search(address, std::regex("http://"))) {
            address  = address.substr(7);
            auto pos = address.find(':');
            host_    = address.substr(0, pos);
            pos == std::string::npos ?
                   port_ = 80 :
                   port_ = std::stoi(address.substr(pos + 1));
        } else {
            auto pos = address.find(':');
            host_    = address.substr(0, pos);
            pos == std::string::npos ?
                   port_ = 80 :
                   port_ = std::stoi(address.substr(pos + 1));
        }

        host_and_port_ = host_ + ":" + std::to_string(port_);
    }

    void BuildHeaders(std::shared_ptr<Request> request) {
        auto& headers = request->headers;
        if (request->method.empty()) request->method = "GET";

        // Host
        if (headers.count("Host") == 0) {
            if (is_ssl_) {
                if (port_ == 443) {
                    headers.emplace(std::make_pair("Host", host_));
                } else {
                    headers.emplace(std::make_pair("Host", host_and_port_));
                }
            } else {
                if (port_ == 80) {
                    headers.emplace(std::make_pair("Host", host_));
                } else {
                    headers.emplace(std::make_pair("Host", host_and_port_));
                }
            }
        }

        if (!keep_alive_) {
            if (headers.count("Connection") == 0) {
                headers.emplace(std::make_pair("Connection", "close"));
            }
        } else {
            if (headers.count("Connection") == 0) {
                headers.emplace(std::make_pair("Connection", "Keep-Alive"));
            }
        }

        // Accept
        if (headers.count("Accept") == 0) { headers.emplace(std::make_pair("Accept", "*/*")); }

        // User-Agent
        if (headers.count("User-Agent") == 0) { headers.emplace(std::make_pair("User-Agent", "leak-http/0.1")); }

        if (request->body.size() > 0) {
            // Content-Type
            if (headers.count("Content-Type") == 0) {
                headers.emplace("Content-Type", "text/plain");
            }

            // Content-Length
            if (headers.count("Content-Length") == 0) {
                headers.emplace(std::make_pair("Content-Length", std::to_string(request->body.size())));
            }
        } else {
            if (request->method == "POST" || request->method == "PUT" ||
                request->method == "PATCH") {
                if (headers.count("Content-Length") == 0) {
                    headers.emplace("Content-Length", "0");
                }
            }
        }
    }

    void SocketWrite(std::shared_ptr<tcp::Socket> socket, std::shared_ptr<Request> request) {
        auto self = shared_from_this();
        auto task = tasks_.front();

        socket->set_onerror([this, self, task](std::string_view err) {
            this->socket_.reset();
            if (!task->complete) {
                AsyncComplete(err);
            }
        });

        socket->PrepareWrite(request->method + " " + request->path + " HTTP/1.1\r\n");

        for (auto& header : request->headers) {
            socket->PrepareWrite(header.first + ": " + header.second + http::crlf);
        }
        socket->PrepareWrite(http::crlf);
        socket->PrepareWrite(request->body);
        socket->Flush();
    }

    void DoRequest() {
        if (intask_ || tasks_.size() == 0) return;
        intask_ = true;

        auto request = tasks_.front()->request;
        auto socket  = socket_.lock();
        // if keep alive socket, reuse it
        if (socket) {
            SocketWrite(std::move(socket), std::move(request));
        } else {
            DoConnect(request);
        }
    }

    void DoConnect(std::shared_ptr<Request> request) {
        auto self = shared_from_this();
        net::tcp::Connect(
            context_,
            host_,
            port_,
            [self, this, request](std::shared_ptr<net::tcp::Socket> socket) {
                SocketWrite(socket, request);
                auto parser = std::make_shared<http::ResponseParser>();
                socket->StartRead([this, self, parser](std::shared_ptr<tcp::Socket> socket, std::string_view data) {
                    auto succ = parser->Parse(data);
                    if (!succ) {
                        socket->close();
                        AsyncComplete("parse response failed");
                        return -1;
                    } else if (parser->IsFinish()) {
                        auto& response = parser->response();

                        if (keep_alive_) socket_ = socket;

                        asio::post(context_, [this, self, socket, response]() {
                            AsyncComplete(std::nullopt, &response);
                            if (!self->keep_alive_) asio::post(context_, [socket]() { socket->close(); });
                        });
                        return -1;
                    }
                    return 0;
                });
            },
            [self, this](std::string_view err) {
                AsyncComplete(err);
            });
    }

    void AsyncComplete(std::optional<std::string_view> err, const Response* rsp = nullptr) {
        assert(tasks_.size() > 0);
        auto task = tasks_.front();
        if (!task->complete) {
            if (err && task->error_handler)
                task->error_handler(*err);
            else if (!err && task->response_handler)
                task->response_handler(*rsp);
            task->complete = true;
        }

        tasks_.pop_front();
        intask_ = false;

        DoRequest();
    }

    asio::io_context&          context_;
    bool                       is_ssl_;
    std::string                host_and_port_;
    std::string                host_;
    uint16_t                   port_;
    bool                       keep_alive_;
    std::weak_ptr<tcp::Socket> socket_;

    std::list<std::shared_ptr<Task>> tasks_;
    bool                             intask_;
};

}  // namespace leak::net::http
